# load_models.py
import os

import torch
import pickle

from app.services.prediction_parameters.model_utils import LSTMModel


class ModelManager:
    def __init__(self):
        self.models = {}
        self.scalers = {}

    def load_model_and_scaler(self, model_path, scaler_path):
        model = LSTMModel()
        model.load_state_dict(torch.load(model_path))
        model.eval()
        with open(scaler_path, 'rb') as f:
            scaler = pickle.load(f)
        return model, scaler

    def load_all(self, base_dir):
        prediction_parameters_dir = os.path.join(base_dir, 'app', 'services', 'prediction_parameters', 'saved_models')

        # 使用 os.path.join 正确拼接文件路径
        self.models['load'], self.scalers['load'] = self.load_model_and_scaler(
            os.path.join(prediction_parameters_dir, "load_kW_lstm_model.pt"),
            os.path.join(prediction_parameters_dir, "load_kW_scaler.pkl")
        )

        self.models['grid'], self.scalers['grid'] = self.load_model_and_scaler(
            os.path.join(prediction_parameters_dir, "grid_price_lstm_model.pt"),
            os.path.join(prediction_parameters_dir, "grid_price_scaler.pkl")
        )

        self.models['wind'], self.scalers['wind'] = self.load_model_and_scaler(
            os.path.join(prediction_parameters_dir, "wind_price_lstm_model.pt"),
            os.path.join(prediction_parameters_dir, "wind_price_scaler.pkl")
        )

        self.models['solar'], self.scalers['solar'] = self.load_model_and_scaler(
            os.path.join(prediction_parameters_dir, "solar_price_lstm_model.pt"),
            os.path.join(prediction_parameters_dir, "solar_price_scaler.pkl")
        )
model_manager = ModelManager()
